#include "coarsen.hpp"
#include "coarsen_heap.hpp"
#include "coarsen_edge.hpp"

template<typename idx_t, typename data_t, typename calc_t>
CoarsenWrapper<idx_t, data_t, calc_t>* coarsen(const MyParCSRMatrix<idx_t, data_t, calc_t> & mat,
    const idx_t use_coord, const calc_t threshold, const idx_t num, const calc_t alpha)
{
    // std::vector<double> record;
    // record.emplace_back(MPI_Wtime());
    const idx_t num_fine = mat.end_row - mat.beg_row;// 本进程在细网格所负责的行数

    CoarsenWrapper<idx_t, data_t, calc_t> * _last = nullptr;
    idx_t * my_map_f2c = new idx_t [num_fine];// 本进程内从细点到粗点的映射关系
    idx_t * rstr_rpt = nullptr;// 限制矩阵的行偏移
    idx_t * rstr_cid = nullptr;// 限制矩阵的列索引
    // 首先初始化，构造一个单位映射
    #pragma omp parallel for schedule(static)
    for (idx_t j = 0; j < num_fine; j++)
        my_map_f2c[j] = j;

    for (idx_t it = 0; it < num; it++) {
        // record.emplace_back(MPI_Wtime());

        const MyParCSRMatrix<idx_t, data_t, calc_t> * fmat_ptr = nullptr;// 细网格矩阵的指针
        CoarsenWrapper<idx_t, data_t, calc_t> * to_free = nullptr;
        if (_last) {
            assert(_last->loc_map_f2c == nullptr);
            assert(_last->R_ci == nullptr);
            assert(_last->R_rp == nullptr);
            fmat_ptr = _last->coar_mat;
            to_free = _last;
        } else {
            fmat_ptr = (& mat);
            to_free = nullptr;
        }

        if (use_coord == 1) {// 使用基于点坐标的成对粗化
            _last = pairwise_coarsen<idx_t, data_t, calc_t>(*fmat_ptr, alpha);
        } else {// 否则，使用基于最小堆的成对粗化
            _last = pairwise_coarsen<idx_t, data_t, calc_t>(*fmat_ptr, threshold, alpha);
        }
        // record.emplace_back(MPI_Wtime());

        // 更新细到粗的映射（用于插值算子）
        #pragma omp parallel for schedule(static)
        for (idx_t j = 0; j < num_fine; j++)
            my_map_f2c[j] = _last->loc_map_f2c[my_map_f2c[j]];
        delete [] _last->loc_map_f2c; _last->loc_map_f2c = nullptr;
        
        // 更新粗到细的映射（用于限制算子）
        if (rstr_rpt == nullptr) { assert(rstr_cid == nullptr);// 转移所有权
            rstr_rpt = _last->R_rp; _last->R_rp = nullptr;
            rstr_cid = _last->R_ci; _last->R_ci = nullptr;
        } else {// 需要合并两次的限制矩阵
            const idx_t num_coar = _last->coar_mat->end_row - _last->coar_mat->beg_row;
            idx_t * merged_rpt = new idx_t [num_coar + 1];
            idx_t * merged_cid = new idx_t [num_fine];
            // 先归纳行偏移
            merged_rpt[0] = 0;

            idx_t thread_flags [omp_get_max_threads()];// 多线程并行的标志位，用来表达该线程当前进行到哪个阶段
            idx_t thread_shared[omp_get_max_threads()];// 多线程相互共享的数据区
            for (idx_t t = 0; t < omp_get_max_threads(); t++) thread_flags[t] = -1;// 阶段编号初始化为-1
            #pragma omp parallel
            {
                int tid = omp_get_thread_num();
                int nt  = omp_get_num_threads();
                int tbeg = 0, tend = 0;
                thread_load_balance(num_coar, nt, tid, tbeg, tend);// 将各线程的负载调整均衡

                idx_t * tl_rpt = new idx_t [tend - tbeg + 1];// 线程内部的行偏移
                tl_rpt[0] = 0;
                for (idx_t topci = tbeg; topci < tend; topci ++) {// 遍历本线程负责的每一行
                    idx_t tot_nnz = 0;
                    // 遍历本行内的非零元，查看它们会对合并后的限制矩阵贡献多少个非零元
                    for (idx_t topcp = _last->R_rp[topci]; topcp < _last->R_rp[topci + 1]; topcp++) {
                        const idx_t midfi = _last->R_ci[topcp];
                        tot_nnz += rstr_rpt[midfi + 1] - rstr_rpt[midfi];
                    }
                    tl_rpt[topci - tbeg + 1] = tl_rpt[topci - tbeg] + tot_nnz;// 线程内局部前缀和
                }
                
                while (tid > 0 && __atomic_load_n( & thread_flags[tid - 1], __ATOMIC_ACQUIRE) < 1) {  }// 等待上一个线程完成
                idx_t prev_tsum = (tid > 0) ? thread_shared[tid - 1] : 0;
                thread_shared[tid] = prev_tsum + tl_rpt[tend - tbeg];// 本线程内的求和存入多线程共享区
                __atomic_store_n( & thread_flags[tid], 1, __ATOMIC_RELEASE);
                
                for (idx_t ci = tbeg; ci < tend; ci++) {// 填充本线程负责范围内的行偏移
                    merged_rpt[ci + 1] = prev_tsum + tl_rpt[ci - tbeg + 1];
                }
                __atomic_store_n( & thread_flags[tid], 2, __ATOMIC_RELEASE);// 保证下一个线程执行时本线程的rpt数据已经写回到内存中
                
                while (tid > 0 && __atomic_load_n( & thread_flags[tid - 1], __ATOMIC_ACQUIRE) < 2) {  }// 等待上一个线程完成
                for (idx_t topci = tbeg; topci < tend; topci ++) {
                    idx_t dst_off = merged_rpt[topci];
                    for (idx_t topcp = _last->R_rp[topci]; topcp < _last->R_rp[topci + 1]; topcp++) {
                        const idx_t midfi = _last->R_ci[topcp];
                        const idx_t src_off = rstr_rpt[midfi];
                        const idx_t num = rstr_rpt[midfi + 1] - src_off;
                        for (idx_t k = 0; k < num; k++)
                            merged_cid[dst_off + k] = rstr_cid[src_off + k];
                        dst_off += num;
                    } assert(dst_off == merged_rpt[topci + 1]);
                    // 排序保证限制矩阵一行内的非零元按列号升序排列
                    bubble_sort<idx_t, idx_t>(merged_rpt[topci + 1] - merged_rpt[topci], merged_cid + merged_rpt[topci]);
                    // qsort(merged_cid + merged_rpt[topci], merged_rpt[topci + 1] - merged_rpt[topci], sizeof(idx_t), cmpfunc<idx_t>);
                }
                delete [] tl_rpt; tl_rpt = nullptr;
            }// implicit sync
            
            delete [] rstr_rpt; rstr_rpt = merged_rpt;
            delete [] rstr_cid; rstr_cid = merged_cid;
            delete [] _last->R_rp; _last->R_rp = nullptr;
            delete [] _last->R_ci; _last->R_ci = nullptr;
        }

        if (to_free != nullptr) { delete to_free; to_free = nullptr; }// 上一层的矩阵，中间结果可以释放掉
    }// num loop

    // 构造一个返回的结构体
    assert(_last);
    assert(_last->loc_map_f2c == nullptr); _last->loc_map_f2c = my_map_f2c;
    assert(_last->R_rp == nullptr); _last->R_rp = rstr_rpt;
    assert(_last->R_ci == nullptr); _last->R_ci = rstr_cid;
    assert(_last->coar_mat);
    return _last;
}
template CoarsenWrapper<int, double, double>* coarsen(const MyParCSRMatrix<int, double, double> & mat,
    const int  use_coord, const double threshold, const int num, const double alpha);
template CoarsenWrapper<int, float , double>* coarsen(const MyParCSRMatrix<int, float , double> & mat,
    const int  use_coord, const double threshold, const int num, const double alpha);
template CoarsenWrapper<int, float , float >* coarsen(const MyParCSRMatrix<int, float , float > & mat,
    const int  use_coord, const float  threshold, const int num, const float  alpha);
